# ægraph (aegraph, or acyclic e-graph) implementation.
An aegraph is a form of e-graph. We will first describe the
e-graph, then the aegraph as a slightly less powerful but highly
optimized variant of it.
The main goal of this library is to be explicitly memory-efficient
and light on allocations. We need to be as fast and as small as
possible in order to minimize impact on compile time in a
production compiler.
## The e-graph
An e-graph, or equivalence graph, is a kind of node-based
intermediate representation (IR) data structure that consists of
*eclasses* and *enodes*. An eclass contains one or more enodes;
semantically an eclass is like a value, and an enode is one way to
compute that value. If several enodes are in one eclass, the data
structure is asserting that any of these enodes, if evaluated,
would produce the value.
An e-graph also contains a deduplicating hash-map of nodes, so if
the user creates the same e-node more than once, they get the same
e-class ID.
In the usual use-case, an e-graph is used to build a sea-of-nodes
IR for a function body or other expression-based code, and then
*rewrite rules* are applied to the e-graph. Each rewrite
potentially introduces a new e-node that is equivalent to an
existing e-node, and then unions the two e-nodes' classes
together.
In the trivial case this results in an e-class containing a series
of e-nodes that are newly added -- all known forms of an
expression -- but Note how if a rewrite rule rewrites into an
existing e-node (discovered via deduplication), rewriting can
result in unioning of two e-classes that have existed for some
time.
An e-graph's enodes refer to *classes* for their arguments, rather
than other nodes directly. This is key to the ability of an
e-graph to canonicalize: when two e-classes that are already used
as arguments by other e-nodes are unioned, all e-nodes that refer
to those e-classes are themselves re-canonicalized. This can
result in "cascading" unioning of eclasses, in a process that
discovers the transitive implications of all individual
equalities. This process is known as "equality saturation".
## The acyclic e-graph (aegraph)
An e-graph is powerful, but it can also be expensive to build and
saturate: there are often many different forms an expression can
take (because many different rewrites are possible), and cascading
canonicalization requires heavyweight data structure bookkeeping
that is expensive to maintain.
This crate introduces the aegraph: an acyclic e-graph. This data
structure stores an e-class as an *immutable persistent data
structure*. An id can refer to some *level* of an eclass: a
snapshot of the nodes in the eclass at one point in time. The
nodes referred to by this id never change, though the eclass may
grow later.
A *union* is also an operation that creates a new eclass id: the
original eclass IDs refer to the original eclass contents, while
the id resulting from the `union()` operation refers to an eclass
that has all nodes.
In order to allow for adequate canonicalization, an enode normally
stores the *latest* eclass id for each argument, but computes
hashes and equality using a *canonical* eclass id. We define such
a canonical id with a union-find data structure, just as for a
traditional e-graph. It is normally the lowest id referring to
part of the eclass.
The persistent/immutable nature of this data structure yields one
extremely important property: it is acyclic! This simplifies
operation greatly:
- When "elaborating" out of the e-graph back to linearized code,
so that we can generate machine code, we do not need to break
cycles. A given enode cannot indirectly refer back to itself.
- When applying rewrite rules, the nodes visible from a given id
for an eclass never change. This means that we only need to
apply rewrite rules at that node id *once*.
## Data Structure and Example
Each eclass id refers to a table entry ("eclass node", which is
different than an "enode") that can be one of:
- A single enode;
- An enode and an earlier eclass id it is appended to (a "child"
eclass node);
- A "union node" with two earlier eclass ids.
Building the aegraph consists solely of adding new entries to the
end of this table of eclass nodes. An enode referenced from any
given eclass node can only refer to earlier eclass ids.
For example, consider the following eclass table:
```plain
eclass/enode table
eclass1 iconst(1)
eclass2 blockparam(block0, 0)
eclass3 iadd(eclass1, eclass2)
```
This represents the expression `iadd(blockparam(block0, 0),
iconst(1))` (as the sole enode for eclass3).
Now, say that as we further build the function body, we add
another enode `iadd(eclass3, iconst(1))`. The `iconst(1)` will be
deduplicated to `eclass1`, and the toplevel `iadd` will become its
own new eclass (`eclass4`).
```plain
eclass4 iadd(eclass3, eclass1)
```
Now we apply our body of rewrite rules, and these results can
combine `x + 1 + 1` into `x + 2`; so we get:
```plain
eclass5 iconst(2)
eclass6 union(iadd(eclass2, eclass5), eclass4)
```
Note that we added the nodes for the new expression, and then we
union'd it with the earlier `eclass4`. Logically this represents a
single eclass that contains two nodes -- the `x + 1 + 1` and `x +
2` representations -- and the *latest* id for the eclass,
`eclass6`, can reach all nodes in the eclass (here the node stored
in `eclass6` and the earlier one in `elcass4`).
## aegraph vs. egraph
Where does an aegraph fall short of an e-graph -- or in other
words, why maintain the data structures to allow for full
(re)canonicalization at all, with e.g. parent pointers to
recursively update parents?
This question deserves further study, but right now, it appears
that the difference is limited to a case like the following:
- expression E1 is interned into the aegraph.
- expression E2 is interned into the aegraph. It uses E1 as an
argument to one or more operators, and so refers to the
(currently) latest id for E1.
- expression E3 is interned into the aegraph. A rewrite rule fires
that unions E3 with E1.
In an e-graph, the last action would trigger a re-canonicalization
of all "parents" (users) of E1; so E2 would be re-canonicalized
using an id that represents the union of E1 and E3. At
code-generation time, E2 could choose to use a value computed by
either E1's or E3's operator. In an aegraph, this is not the case:
E2's e-class and e-nodes are immutable once created, so E2 refers
only to E1's representation of the value (a "slice" of the whole
e-class).
While at first this sounds quite limiting, there actually appears
to be a nice mutually-beneficial interaction with the immediate
application of rewrite rules: by applying all rewrites we know
about right when E1 is interned, E2 can refer to the best version
when it is created. The above scenario only leads to a missed
optimization if:
- a rewrite rule exists from E3 to E1, but not E1 to E3; and
- E3 is *cheaper* than E1.
Or in other words, this only matters if there is a rewrite rule
that rewrites into a more expensive direction. This is unlikely
for the sorts of rewrite rules we plan to write; it may matter
more if many possible equalities are expressed, such as
associativity, commutativity, etc.
Note that the above represents the best of our understanding, but
there may be cases we have missed; a more complete examination of
this question would involve building a full equality saturation
loop on top of the (a)egraph in this crate, and testing with many
benchmarks to see if it makes any difference.
## Rewrite Rules (FLAX: Fast Localized Aegraph eXpansion)
The most common use of an e-graph or aegraph is to serve as the IR
for a compiler. In this use-case, we usually wish to transform the
program using a body of rewrite rules that represent valid
transformations (equivalent and hopefully simpler ways of
computing results). An aegraph supports applying rules in a fairly
straightforward way: whenever a new eclass entry is added to the
table, we invoke a toplevel "apply all rewrite rules" entry
point. This entry point creates new nodes as needed, and when
done, unions the rewritten nodes with the original. We thus
*immediately* expand a new value into all of its representations.
This immediate expansion stands in contrast to a traditional
"equality saturation" e-egraph system, in which it is usually best
to apply rules in batches and then fix up the
canonicalization. This approach was introduced in the `egg`
e-graph engine [^1]. We call our system FLAX (because flax is an
alternative to egg): Fast Localized Aegraph eXpansion.
The reason that this is possible in an aegraph but not
(efficiently, at least) in a traditional e-graph is that the data
structure nodes are immutable once created: an eclass id will
always refer to a fixed set of enodes. There is no
recanonicalizing of eclass arguments as they union; but also this
is not usually necessary, because args will have already been
processed and eagerly rewritten as well. In other words, eager
rewriting and the immutable data structure mutually allow each
other to be practical; both work together.
[^1]: M Willsey, C Nandi, Y R Wang, O Flatt, Z Tatlock, P
Panchekha. "egg: Fast and Flexible Equality Saturation." In
POPL 2021.